#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
功能超级完善的点云可视化工具 (The Ultimate Point Cloud Visualizer)
作者: [你的名字或AI助手名称]
版本: 2.2 (已修复裁剪功能 bug)
------------------------------------------------------------------------------------
本脚本是一个多功能的命令行工具，旨在满足各种点云可视化需求。

核心功能:
1.  支持单个或多个点云文件的同时加载和可视化。
2.  提供多种着色方案，包括默认颜色、高度着色、法向量着色和统一着色。
3.  可自定义渲染选项，如点大小、背景色、坐标轴显示。
4.  支持EDL（Eye-Dome Lighting）高级着色，增强深度感知。
5.  提供交互式连续裁剪模式，用于从高密度点云中分离目标区域。
6.  支持保存当前视图为图像文件。
7.  支持保存和加载相机视角，用于复现可视化场景。

依赖库:
- open3d
- numpy
- matplotlib

使用示例:
# 1. 基本可视化
python visualize_pcd.py cloud1.pcd

# 2. 同时可视化两个点云，并使用统一颜色区分
python visualize_pcd.py cloud1.ply cloud2.xyz --color uniform

# 3. 使用高度着色，并显示坐标轴
python visualize_pcd.py terrain.pcd --color z_axis --show_axes

# 4. 使用EDL着色，并保存为图片
python visualize_pcd.py model.ply --edl --save_image result.png

# 5. 进入交互式连续裁剪模式
python visualize_pcd.py dense_scene.pcd --pick_points

# 6. 保存当前视角，下次加载
#    第一次: python visualize_pcd.py model.ply --save_view my_view.json
#    第二次: python visualize_pcd.py model.ply --load_view my_view.json
"""

import open3d as o3d
import numpy as np
import argparse
import os
import json
from matplotlib import cm

def get_color_map(name='viridis'):
    """获取一个matplotlib的颜色映射"""
    #
    # cmap in older versions of matplotlib needed to be imported from 'pylab'
    try:
        from matplotlib import cm
        return cm.get_cmap(name)
    except ImportError:
        from matplotlib import pylab
        return pylab.get_cmap(name)


def color_pcd_by_axis(pcd, axis='z', cmap_name='viridis'):
    """根据指定轴向对点云进行着色"""
    points = np.asarray(pcd.points)
    if axis == 'x':
        values = points[:, 0]
    elif axis == 'y':
        values = points[:, 1]
    else: # 'z'
        values = points[:, 2]

    # 归一化到 [0, 1]
    normalized_values = (values - np.min(values)) / (np.max(values) - np.min(values))
    
    # 获取颜色映射
    cmap = get_color_map(cmap_name)
    colors = cmap(normalized_values)[:, :3]  # cmap返回RGBA, 我们只需要RGB
    
    pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd

def main():
    parser = argparse.ArgumentParser(
        description="一个功能强大的点云可视化工具。",
        formatter_class=argparse.RawTextHelpFormatter
    )
    
    parser.add_argument(
        "input_files", nargs='+',
        help="一个或多个输入的点云文件路径。"
    )

    # --- 可视化风格组 ---
    style_group = parser.add_argument_group('可视化风格 (Visualization Style)')
    style_group.add_argument(
        "--color", type=str, default='default',
        choices=['default', 'z_axis', 'uniform', 'normal'],
        help="""设置点云着色方案:
  default - 使用文件自带的颜色，如果没有则为灰色。
  z_axis  - 根据Z轴高度进行着色 (热力图)。
  uniform - 为每个点云指定一个统一的颜色。
  normal  - 根据法向量方向着色 (需先计算法线)。"""
    )
    style_group.add_argument('--point_size', type=float, default=2.0, help="设置点的大小。")
    style_group.add_argument(
        '--bg_color', type=str, default='dark', choices=['dark', 'light', 'gray'],
        help="设置背景颜色。"
    )

    # --- 渲染特性组 ---
    feature_group = parser.add_argument_group('渲染特性 (Rendering Features)')
    feature_group.add_argument('--show_axes', action='store_true', help="显示世界坐标系XYZ轴。")
    feature_group.add_argument('--edl', action='store_true', help="启用Eye-Dome Lighting着色，增强深度感。")

    # --- 交互与输出组 ---
    io_group = parser.add_argument_group('交互与输出 (Interaction & Output)')
    io_group.add_argument(
        '--pick_points', action='store_true',
        help="""进入交互式连续裁剪模式:
  - 可以通过多次裁剪来逐步分离出高密度点云中的目标。
  - 在窗口中按 [K] 键锁定屏幕，然后使用 [Ctrl/Cmd + 鼠标左键拖动] 或 
    [Shift + 鼠标左键拖动] 来选择区域。
  - 选择完毕后，关闭窗口，裁剪出的点云会替换原来的点云并重新显示，
    以进行下一次更精细的裁剪。
  - 再次关闭窗口即可结束并输出最终选中的点信息。"""
    )
    io_group.add_argument('--save_image', type=str, default=None, help="将最终视图保存为图像文件 (例如: view.png)。")
    io_group.add_argument('--save_view', type=str, default=None, help="将当前相机视角参数保存到JSON文件。")
    io_group.add_argument('--load_view', type=str, default=None, help="从JSON文件加载并应用相机视角。")

    args = parser.parse_args()

    # --- 1. 加载点云 ---
    pcds = []
    print("INFO: 正在加载点云文件...")
    for f in args.input_files:
        if not os.path.exists(f):
            print(f"WARNING: 文件 '{f}' 不存在，已跳过。")
            continue
        pcd = o3d.io.read_point_cloud(f)
        if not pcd.has_points():
            print(f"WARNING: 文件 '{f}' 为空，已跳过。")
            continue
        pcds.append(pcd)
        print(f"  - 已加载 '{f}' (包含 {len(pcd.points)} 个点)。")
    
    if not pcds:
        print("ERROR: 未能加载任何有效的点云文件。脚本终止。")
        return

    # --- 2. 准备要绘制的几何体 ---
    geometries = []
    
    # 预定义的鲜艳颜色循环，用于uniform模式
    color_cycle = [
        [1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 0], [0, 1, 1], [1, 0, 1] 
    ]

    for i, pcd in enumerate(pcds):
        if args.color == 'z_axis':
            pcd = color_pcd_by_axis(pcd, 'z')
        elif args.color == 'uniform':
            color = color_cycle[i % len(color_cycle)]
            pcd.paint_uniform_color(color)
        elif args.color == 'normal':
            if not pcd.has_normals():
                print(f"INFO: 点云 '{args.input_files[i]}' 没有法线，正在估算...")
                pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
            pcd.normalize_normals()
            # 将法向量从[-1, 1]映射到[0, 1]作为颜色
            colors = (np.asarray(pcd.normals) + 1.0) / 2.0
            pcd.colors = o3d.utility.Vector3dVector(colors)

        geometries.append(pcd)
    
    if args.show_axes:
        # 创建一个坐标轴几何体
        coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])
        geometries.append(coord_frame)

    # --- 3. 执行可视化或交互 ---
    if args.pick_points:
        # ########## FIXED SECTION ##########
        print("\n----------- 进入交互式连续裁剪模式 -----------")
        print("说明:")
        print("  1. 在可视化窗口中，调整视角，然后按 [K] 锁定屏幕交互。")
        print("  2. 使用 [Ctrl/Cmd + 鼠标左键] 或 [Shift + 鼠标左键] 拖动以绘制选择框。")
        print("  3. 框选完毕后，按 [Q] 或关闭窗口。")
        print("  4. 被裁剪出的点云将作为新的目标，并自动打开新窗口以进行下一轮裁剪。")
        print("  5. 当您在新窗口中不进行任何裁剪，直接按 [Q] 关闭时，流程结束。")
        print("--------------------------------------------------")

        # 合并所有点云以便于统一裁剪
        pcd_to_crop = o3d.geometry.PointCloud()
        for geom in geometries:
            if isinstance(geom, o3d.geometry.PointCloud):
                pcd_to_crop += geom
        
        if not pcd_to_crop.has_points():
            print("ERROR: 没有可用于裁剪的点云。")
            return

        final_selection = None
        while True:
            print(f"\nINFO: 打开裁剪窗口... 当前点数: {len(pcd_to_crop.points)}")
            print("INFO: 请进行裁剪，或直接按 [Q] 结束并使用当前点云作为最终选择。")
            
            # 使用 VisualizerWithEditing 进行交互式裁剪
            # 注意：旧版 open3d 可能使用 pick_points 方法，新版使用 crop_geometry
            # 此处使用 crop_geometry 的标准方式
            vis = o3d.visualization.VisualizerWithEditing()
            vis.create_window("交互式裁剪 (按 'K' 锁定屏幕并选择，按 'Q' 确认)", width=1280, height=720)
            vis.add_geometry(pcd_to_crop)
            vis.run()  # 用户在此进行交互
            vis.destroy_window()

            # *** 关键修正点 ***
            # 使用 get_cropped_geometry() 获取裁剪后的新几何体
            cropped_pcd = vis.get_cropped_geometry()

            if not cropped_pcd or not cropped_pcd.has_points():
                # 用户没有进行裁剪就关闭了窗口，意味着当前的点云就是最终想要的
                final_selection = pcd_to_crop
                print("INFO: 未进行新的裁剪或裁剪结果为空，流程结束。")
                break
            
            # 将裁剪后的点云作为下一次循环的目标
            pcd_to_crop = cropped_pcd
        
        if final_selection and final_selection.has_points():
            points = np.asarray(final_selection.points)
            print("\n----------- 最终选定的点信息 -----------")
            print(f"总计选定 {len(points)} 个点。")
            # 为防止刷屏，最多显示前100个点的坐标
            for i, point in enumerate(points):
                if i >= 100:
                    print(f"... (及其他 {len(points) - 100} 个点)")
                    break
                print(f"  - 点 {i+1}: ({point[0]:.4f}, {point[1]:.4f}, {point[2]:.4f})")
            print("------------------------------------------")
        else:
            print("\nINFO: 没有选定任何点。")
        # ########## END OF FIXED SECTION ##########

    else:
        # 标准可视化流程
        bg_colors = {'dark': [0.1, 0.1, 0.1], 'light': [1.0, 1.0, 1.0], 'gray': [0.8, 0.8, 0.8]}
        
        def render_and_save(geoms, args):
            vis = o3d.visualization.Visualizer()
            vis.create_window("点云可视化工具", width=1280, height=720)
            
            opt = vis.get_render_option()
            opt.point_size = args.point_size
            opt.background_color = np.asarray(bg_colors[args.bg_color])
            if args.edl:
                opt.point_color_option = o3d.visualization.PointColorOption.EyeDomeLighting

            for geom in geoms:
                vis.add_geometry(geom)

            if args.load_view:
                try:
                    with open(args.load_view, 'r') as f:
                        cam_params = json.load(f)
                    ctr = vis.get_view_control()
                    pinhole_params = ctr.convert_to_pinhole_camera_parameters()
                    pinhole_params.extrinsic = np.array(cam_params['extrinsic'])
                    pinhole_params.intrinsic.intrinsic_matrix = np.array(cam_params['intrinsic_matrix'])
                    ctr.convert_from_pinhole_camera_parameters(pinhole_params, allow_arbitrary=True)
                    print(f"INFO: 已从 '{args.load_view}' 加载相机视角。")
                except Exception as e:
                    print(f"ERROR: 加载视角失败: {e}")
            
            print("\nINFO: 正在打开可视化窗口。按 [Q] 关闭。")
            vis.run() # 主循环

            # 在窗口关闭后执行保存操作
            if args.save_view:
                ctr = vis.get_view_control()
                params = ctr.convert_to_pinhole_camera_parameters()
                cam_params = {
                    'extrinsic': params.extrinsic.tolist(),
                    'intrinsic_matrix': params.intrinsic.intrinsic_matrix.tolist()
                }
                with open(args.save_view, 'w') as f:
                    json.dump(cam_params, f, indent=4)
                print(f"INFO: 相机视角已保存到 '{args.save_view}'。")

            if args.save_image:
                vis.capture_screen_image(args.save_image, do_render=True)
                print(f"INFO: 视图已保存为图像 '{args.save_image}'。")

            vis.destroy_window()

        render_and_save(geometries, args)

    print("\nINFO: 脚本执行完毕。")

if __name__ == '__main__':
    main()